

import tensorflow as tf
import numpy as np
import os
import matplotlib.pyplot as plt
import uuid

from vehicle_detection import *
from tensorflow.core.protobuf import saved_model_pb2
from tensorflow.python.util import compat
from PIL import Image,ImageDraw, ImageFont
from shutil import copy2
from tensorflow.python.platform import gfile

from flask import Flask, request, redirect, send_from_directory, url_for


app = Flask(__name__)


PATH_TO_DETECTION_PB = './ckpt/detection/ssd_mobilenet_v2_coco_2018_03_29.pb'    
PATH_TO_DETECTION_LABELS = './ckpt/detection/mscoco_label_map.pbtxt'   
DETECTION_NUM_CLASSES = 90


PATH_TO_CLASSIFY_PB = './ckpt/classify/freezed.pb'    
PATH_TO_CLASSIFY_LABELS = './ckpt/classify/label.label'   
CLASSIFY_NUM_CLASSES = 764

app._static_folder = 'static'
UPLOAD_FOLDER = 'static/test_image'
OUTPUT_FOLDER = 'static/out_image'
ALLOWED_EXTENSIONS = set(['jpg','JPG', 'jpeg', 'JPEG', 'png'])


def allowed_files(filename):
    return '.' in filename and filename.rsplit('.', 1)[1] in ALLOWED_EXTENSIONS

def rename_filename(old_filename):
    basename = os.path.basename(old_filename)
    name, ext = os.path.splitext(basename)
    new_filename = str(uuid.uuid1()) + ext
    return new_filename  

def vehicle_detection_classify(image_name, detection_config, classify_config):    
   
    detection_graph = tf.Graph()
    #define a new graph
    with detection_graph.as_default():      
        with gfile.FastGFile(detection_config, 'rb') as f: #加载模型
            graph_def = tf.GraphDef()
            graph_def.ParseFromString(f.read())
            tf.import_graph_def(graph_def, name='')  # 导入计算图

        with tf.Session(graph=detection_graph) as sess:
            image_tensor = detection_graph.get_tensor_by_name('image_tensor:0')     
            detection_boxes = detection_graph.get_tensor_by_name('detection_boxes:0')  
            detection_scores = detection_graph.get_tensor_by_name('detection_scores:0')   
            detection_classes = detection_graph.get_tensor_by_name('detection_classes:0') 
            num_detections = detection_graph.get_tensor_by_name('num_detections:0')

            image = Image.open(image_name)
            #open the image
            img_width,img_height = image.size
            image_np = image_to_numpy(image)  
            #transform the image to numpy that can be uesd in tensorflow API
            image_np_expanded = np.expand_dims(image_np, axis=0)     

            (boxes, scores, classes, num) = sess.run([detection_boxes, detection_scores, detection_classes, num_detections],\
                                                      feed_dict={image_tensor: image_np_expanded})
            #The above core comes from research/object_detection/inference.py
            
    vehicle_name,vehicle_box = detection_box_accuracy(boxes, scores, classes)
    ### This Method to get the bounding boxes include car or truck and their scores
    
    if vehicle_name and vehicle_box:
        classify_graph = tf.Graph()
        with classify_graph.as_default():        
            with open(classify_config, 'rb') as f:
                graph_def = tf.GraphDef()
                graph_def.ParseFromString(f.read())
                tf.import_graph_def(graph_def, name='')
            #[print(n.name) for n in tf.get_default_graph().as_graph_def().node]    
           
            image_preprocessed_list = crop_box(vehicle_box, image,image_np)
            ###This Method to get the image of object in the bounding-box and prepare to send to classification model
            
            sess=tf.Session()
            image_preprocessed_list = sess.run(image_preprocessed_list)
            img_list = np.array(image_preprocessed_list)

            img_path_list = []
            for index,im in enumerate(img_list):
                img_path = os.path.join(OUTPUT_FOLDER, os.path.basename(str(index)+'.jpg'))
                img_path_list.append(img_path)
                plt.imsave(img_path, im)

            # stack tensor list to tensor
            with tf.Session(graph=classify_graph) as sess:
                softmax_tensor = sess.graph.get_tensor_by_name('final_probs:0')
                predictions = np.zeros(shape=(len(img_list),CLASSIFY_NUM_CLASSES))
                for index,img in enumerate(img_path_list):
                    image_data = open(img, 'rb').read()
                    pr = sess.run(softmax_tensor, feed_dict={'input:0': image_data})
                    predictions[index] = pr
            # get the predictions from classfication model
            
        vehicle_predict_name,vehicle_predict_box = classification_box_accuracy(predictions,vehicle_box)

        img_width,img_height = image.size
        test_image = Image.fromarray(image_np)                   
        draw = ImageDraw.Draw(test_image)                            
        use_normalized_coordinates=True

        for i in range(len(vehicle_predict_box)):
            ymin = vehicle_predict_box[i][0]
            xmin = vehicle_predict_box[i][1]
            ymax = vehicle_predict_box[i][2]
            xmax = vehicle_predict_box[i][3]

            if use_normalized_coordinates:
                (left,right,top,bottom) = (xmin * img_width, xmax * img_width,
                                           ymin * img_height, ymax * img_height)
            else:
                (left,right,top,bottom) = (xmin,xmax,ymin,ymax)
            draw.line([(left, top), (left, bottom), (right, bottom),(right, top),(left,top)], width=8,fill='cyan')   

            try:
                font = ImageFont.truetype('./simhei.ttf',35,encoding='utf-8')    
            except IOError:
                font = ImageFont.load_default()

            text_width, text_height = font.getsize(vehicle_predict_name[i])  
            text_bottom = top

            margin = np.ceil(0.05 * text_height)
            draw.rectangle([(left, text_bottom - text_height - 2 * margin), (left + text_width,text_bottom)],fill='cyan')
            draw.text( (left + margin, text_bottom - text_height - margin),
                        vehicle_predict_name[i],
                        fill='black',
                        font=font)

        im = np.array(test_image)    
        plt.imsave(os.path.join(OUTPUT_FOLDER, os.path.basename(image_name)), im)

        image_height = int(img_height/2)
        image_width = int(img_width/2)

        image_detection = OUTPUT_FOLDER + '/%s' % os.path.basename(os.path.join(OUTPUT_FOLDER, image_name))
        image_tag = '<img src="%s" height="%d" width="%d"></img><p>'
        image_detection_tag = image_tag % (image_detection,image_height,image_width) 

        show_result = '<b>检测到的车辆型号如下：</b><br/>'
        for name in vehicle_predict_name:
            show_result += name + '<br>'
        show_all_result  = image_detection_tag + show_result + '<br>'
        return show_all_result
    
   
    elif not vehicle_name:
        plt.imsave(os.path.join(OUTPUT_FOLDER, os.path.basename(image_name)),image_np)
        
        image_height = int(img_height/2)
        image_width = int(img_width/2)

        image_detection = OUTPUT_FOLDER + '/%s' % os.path.basename(os.path.join(OUTPUT_FOLDER, image_name))
        image_tag = '<img src="%s" height="%d" width="%d"></img><p>'
        image_detection_tag = image_tag % (image_detection,image_height,image_width) 

        show_result = '<b>图片中没有汽车</b><br/>'
        show_all_result  = image_detection_tag + show_result + '<br>'
        return show_all_result
        
    
@app.route("/", methods=['GET','POST'])
def root(): 
    vehicle_result = """
        <!doctype html>
       
        <title>车辆检测及型号识别</title>
        <font size=5 color=black> 车辆检测及型号识别 </font><br><br> 
        <font size=4 color=blue> 上传图片，检测图片中车辆位置、并识别其型号 </font> <br>               
        <form action="" method=post enctype=multipart/form-data>
        <p><input type=file name=file value='选择图片' style='font-size:20px'> 
            <input type=submit value='上传图片' style='font-size:20px'> 
        </form>
        <p>%s</p>
        """ % "<br>"  
    
    if request.method == 'POST':
        file = request.files['file']
        old_filename = file.filename
        if file and allowed_files(old_filename):
            filename = rename_filename(old_filename)
            image_path = os.path.join(UPLOAD_FOLDER, filename)
            file.save(image_path)
            type_name = 'N/A' 
            out_image = vehicle_detection_classify(image_path,detection_config=PATH_TO_DETECTION_PB,classify_config=PATH_TO_CLASSIFY_PB)     
            return vehicle_result + out_image 
    return vehicle_result
       
    
if __name__ == "__main__":
    app.run(host='127.0.0.1', port=5000, debug=True, threaded=True)
     
